#include "CudaModule.h"

using namespace OBR;

std::map<std::string, std::shared_ptr<CudaModule>> CudaModule::mapping;

CudaModule::CudaModule(std::string _cuda_rel_source_path) : module_compile_options({})
{
    cuda_source_path = std::string(PROJ_ROOT) + "/" + _cuda_rel_source_path;
    this->module_compile_options.boundValues = 0;
    this->module_compile_options.debugLevel = OPTIX_COMPILE_DEBUG_LEVEL_DEFAULT;
    this->module_compile_options.optLevel = OPTIX_COMPILE_OPTIMIZATION_DEFAULT;
    this->module_compile_options.maxRegisterCount = OPTIX_COMPILE_DEFAULT_MAX_REGISTER_COUNT;
}

CudaModule::~CudaModule()
{
    OPTIX_CHECK_NOEXCEPT(optixModuleDestroy(this->optix_module));
}

bool CudaModule::compile()
{
    // read cuda source as str
    std::string cuda_source_content;
    std::ifstream cuda_source_file(this->cuda_source_path.c_str(), std::ios::binary);
    if (cuda_source_file.good())
    {
        std::vector<char> buffer =
            std::vector<char>(std::istreambuf_iterator<char>(cuda_source_file), {});
        cuda_source_content.assign(buffer.begin(), buffer.end());
    }
    else
    {
        std::cerr << "Couldn't open source file @ '" + this->cuda_source_path + "'" << std::endl;
        throw std::runtime_error("error opening source file");
    }

    // compile cuda source as ptx
    nvrtcProgram prog = 0;
    nvrtcCreateProgram(&prog, cuda_source_content.c_str(), this->cuda_source_path.c_str(), 0, NULL,
                       NULL);
    // nvrtc options
    std::vector<std::string> options_str;
    std::vector<const char *> options;
    // add include dirs
    std::string proj_shader_inc_dirs_raw = PROJ_SHADER_INCLUDE_DIRS;
    std::vector<std::string> shader_include_dirs =
        CompileConfig::parseProjIncludeDirs(proj_shader_inc_dirs_raw);
    for (std::string &dir : shader_include_dirs)
    {
        std::string opt = "-I" + dir;
        options_str.push_back(opt);
    }
    for (std::string &opt : options_str)
    {
        options.push_back(opt.c_str());
    }
    // TODO: fill compiler options
    const char *compiler_options[] = {CUDA_NVRTC_OPTIONS};
    std::copy(std::begin(compiler_options), std::end(compiler_options),
              std::back_inserter(options));
    // JIT compile
    std::string nvrtc_log;
    const nvrtcResult compile_result =
        nvrtcCompileProgram(prog, (int)options.size(), options.data());
    size_t log_size = 0;
    nvrtcGetProgramLogSize(prog, &log_size);
    nvrtc_log.resize(log_size);
    if (log_size > 1)
    {
        nvrtcGetProgramLog(prog, &nvrtc_log[0]);
        // TODO: do something with log
    }
    if (compile_result != NVRTC_SUCCESS)
    {
        std::cerr << "Error: NVRTC failed." << std::endl;
        std::cerr << nvrtc_log << std::endl;
        return false;
    }

    // Retrieve PTX code
    size_t ptx_size = 0;
    nvrtcGetPTXSize(prog, &ptx_size);
    this->compiled_ptx.resize(ptx_size);
    nvrtcGetPTX(prog, &this->compiled_ptx[0]);

    // Cleanup nvrtc
    nvrtcDestroyProgram(&prog);

    return true;
}

bool CudaModule::link(OptixDeviceContext &ctx,
                      OptixPipelineCompileOptions &pipeline_compile_options)
{
    if (!this->ready)
    {
        if (this->compile())
        {
            char *log = new char[MAX_INTERNAL_LOG_SIZE];
            size_t log_size = MAX_INTERNAL_LOG_SIZE;
            size_t ptx_size = this->compiled_ptx.length();
            // create module
            OPTIX_CHECK(
                optixModuleCreateFromPTX(ctx,                             // context
                                         &(this->module_compile_options), // module compile options
                                         &pipeline_compile_options,       // pipline compile options
                                         this->compiled_ptx.c_str(), ptx_size, // compiled ptx
                                         log, &log_size,                       // log
                                         &(this->optix_module));
                std::cout << "=============== Module Link Log ===============" << std::endl;
                std::cout << "for module @ " << this->cuda_source_path << ":" << std::endl;
                std::cout << log << std::endl; delete[] log;);
            this->ready = true;
        }
        else
        {
            throw std::runtime_error("error compiling cuda source file @ " +
                                     this->cuda_source_path);
        }
    }
    else
    {
        std::cerr << "bad call: you shouldn't link a module twice" << std::endl;
    }

    return this->ready;
}

std::shared_ptr<CudaModule> CudaModule::addModule(std::string _cuda_source_path)
{
    CudaModule::mapping[_cuda_source_path] = std::make_shared<CudaModule>(_cuda_source_path);
    return CudaModule::mapping[_cuda_source_path];
}

std::shared_ptr<CudaModule> CudaModule::getModule(std::string cuda_source_path)
{
    if (CudaModule::mapping.find(cuda_source_path) != CudaModule::mapping.end())
    {
        return CudaModule::mapping[cuda_source_path];
    }
    else
    {
        return CudaModule::addModule(cuda_source_path);
    }
}

void CudaModule::linkAllModules(OptixDeviceContext &ctx,
                                OptixPipelineCompileOptions &pipeline_compile_options)
{
    for (auto const &[key, val] : CudaModule::mapping)
    {
        val->link(ctx, pipeline_compile_options);
    }
}